import numpy as np
from skimage.metrics import structural_similarity,peak_signal_noise_ratio


def batch_PSNR(img, imclean, data_range):
    Img = img.data.cpu().numpy().astype(np.float32)
    Iclean = imclean.data.cpu().numpy().astype(np.float32)
    PSNR = 0
    for i in range(Img.shape[0]):
        PSNR += peak_signal_noise_ratio(Iclean[i,0,:,:], Img[i,0,:,:], data_range=data_range)
    return (PSNR/Img.shape[0])

def batch_SSIM(img, imclean, data_range):
    Img = img.data.cpu().numpy().astype(np.float32)
    Iclean = imclean.data.cpu().numpy().astype(np.float32)
    SSIM = 0
    for i in range(Img.shape[0]):
        SSIM +=structural_similarity(Iclean[i,0,:,:], Img[i,0,:,:],data_range=data_range,win_size=11)
    return (SSIM/Img.shape[0])
